/*
   Features
   1. Type rbnode is actually a pointer to its body. this avoids
   the complexities in passing arguments.
   2. Type rbtree is to maintain some meta data about a tree, as
   they might change during manipulation.
   3. Different from normal rbtree, this tree is 'indexed', allowing
   indexing in logarithm time complexity.
*/

#include <stdlib.h>
#include <stdio.h>

#include "rbtree.h"

#define black  0
#define red    1

#define is_black(x) ((x)==NULL || (x) -> color == black)
#define is_red(x)   ((x)!=NULL && (x) -> color == red)

#define _print(x)								\
    (((x)->color)								\
     ?printf("\033[31m%d %lf %lf\033[0m\n", (x)->idx, (x)->value, (x)->content) \
     :printf("%d %lf %lf\n",                (x)->idx, (x)->value, (x)->content))


void _print_tree(rbnode node, int n){
    if ( node->child[0] != NULL )
	_print_tree(node->child[0], n+1);
    for(int i=0; i<n; i++)printf("  ");
    _print(node);
    if ( node->child[1] != NULL )
	_print_tree(node->child[1], n+1);
}


// prepare for insert and delete function

// Initializer

rbnode new_node(double value, double content){
    rbnode x = (rbnode)malloc(sizeof(struct _rbnode));
    *x = (struct _rbnode){
        .content = content,
        .color   = red,
        .dir     = 0,
        .value   = value,
        .parent  = NULL,
        .child   = {NULL, NULL}
    };
    return x;
}

rbtree new_tree(double value, double content){
    rbtree tree = (rbtree)malloc(sizeof(struct _rbtree));
    tree->root  = new_node(value, content);
    tree->root->color = black;
    tree->root  = (rbnode)malloc(sizeof(struct _rbnode));
    tree->n     = 1;
    return tree;
}

// update function is used in the end of either insertion or deletion

#define update(x) while((x)->root->parent!=NULL)x->root=x->root->parent

// find according to binary search tree
int find(rbtree tree, double value, rbnode* ans){
    rbnode node = tree->root;
    int idx = 0;
    while (1){
	if ( node->value == value ){
	    *ans = node;
	    return found;
	}
	idx = ( node->value > value )? left: right;
	if ( node->child[idx] == NULL ){
	    *ans = node;
	    return idx;
	}
	node = node->child[idx];
    }
}

rbnode get(rbtree tree, int idx){
    if (idx<0 || idx >= tree->n)
	return NULL;
    rbnode x = tree->root;
    int dir;
    while ( x->idx != idx ){
	dir  = (idx < x->idx) ? left : right;
	idx -= x->idx;
	x    = x->child[dir];
    }
    return x;
}

void set_child(rbnode x, rbnode y, int dir){
    x -> child[dir] = y;
    if (y!=NULL){
	y -> parent = x;
	y -> dir    = dir;
    }
}

void rotate(rbnode node, int dir){
    // child must not be NULL
    rbnode child  = node  -> child[!dir];
    rbnode gchild = child -> child[dir];
    if ( gchild != NULL )
	gchild->idx += child->idx;
    child->idx += node->idx;
    node->idx  -= child->idx;
    if ( node->parent == NULL )
	child->parent = NULL;
    else
	set_child(node->parent, child, node->dir);
    set_child(node, gchild, !dir);
    set_child(child, node, dir);
}

// following macros is defined to shorten the expressions 
// in the main loop in insert

#define sibling(x)  ((x)->parent->child[!((x)->dir)])

#define _sibling	(sibling(node))
#define _parent		(node->parent)
#define _g_parent	(_parent->parent)
#define _uncle		(sibling(_parent))
#define _dir		(node->dir)

void insert(rbtree tree, double value, double content){
    rbnode node;
    int flag = find(tree, value, &node);
    if ( flag == found )
	return;
    rbnode _node = new_node(value, content);
    _node->idx = (flag==left) ? 0 : 1;
    set_child(node, _node,  flag);
    for (node=_node; _parent!=NULL; node=_parent)
	if ( _dir == left ){
	    _parent->idx++;
	    node->idx--;
	}
    node = _node;
    // main loop
    while (1){
	if ( _parent == NULL){
	    node->color = black;
	    goto END;
	}

	if ( is_black(_parent) ){                       /* R/B/X/Y: red/black/arbitary */
	    goto END;                                   /* number : The ranking        */
	}                                               /* braket : The current node   */

	if ( is_red(_uncle) ){                          /*      B3         (R3)    */
	    _parent -> color = _uncle -> color = black; /*     /  \        /  \    */
	    _g_parent -> color = red;                   /*    R2  R4 ->   B2  B4   */
	    node = _g_parent;                           /*   /           /         */
	    continue;                                   /* (R1)         R1         */
	}
                                                        /*   B3           B3       */
	if ( _dir != _parent->dir ){                    /*  /  \         /  \      */
	    node = _parent;                             /* R1  B4 ->    R2   B4    */
	    rotate(node, _dir);                         /*  \          /           */
	}                                               /* (R2)      (R1)          */

	_parent->color = black;                         /*      B3         B2      */
	_g_parent->color = red;                         /*     /  \       /  \     */
	rotate(_g_parent, !_dir);                       /*    R2  B4 -> (R1)  R3   */
    }                                                   /*   /                 \   */
END:                                                    /* (R1)                 B4 */
    update(tree);
    tree->n++;
}

#define _child (node->child)

int delete(rbtree tree, double value){
    rbnode node;
    int flag = find(tree, value, &node);
    if ( flag!=found )
	return node_not_exist;
    rbnode _node = node;
    if ( _child[0]!=NULL ){
	node = _child[0];
	while ( _child[1]!=NULL )
	    node = _child[1];
    } else if ( _child[1]!=NULL ){        /*    X3      X3  */
	for(;_parent!=NULL;node=_parent){ /*   /       /    */
	    if(_dir==left){               /* (B1)  -> R2    */
		node->idx++;              /*   \            */
		_parent->idx--;           /*    R2          */
	    }
	}
	node = _node;
	_child[1]->color= black;
	_child[1]->idx  = node->idx;
	set_child(_parent, _child[1], _dir);
	free(node);
	update(tree);
	tree->n--;
	return node_deleted;
    }
    _node->value = node->value;
    _node->content = node->content;
    _node = node;
    for (;_parent!=NULL;node=_parent)
	if (_dir==left){
	    node->idx ++;
	    _parent->idx --;
	}
    node = _node;
    if ( node->color == red )
	goto END;
    // copy color in advance
    node->color = (is_black(node->child[0]))? black: red; 
    while (1){
	if ( node->color == red || _parent == NULL ){
	    node->color = black;
	    goto END;
	}

	if ( is_red(_sibling) ){                              /*    B2              B4    */
	    _sibling->color = black;                          /*   /  \            /  \   */
	    _parent->color = red;                             /* (B1)  R4   ->    R2   B5 */
	    rotate(_parent, _dir);                            /*      /  \       /  \     */
	}                                                     /*     B3   B5  (B1)  B3    */

	if ( is_red(_sibling->child[!_dir]) ){                /*    X2             X4    */
	    _sibling->color = _parent->color;                 /*   /  \           /  \   */
	    _parent->color = black;                           /* (B1)  B4   ->   B2  B5  */
	    _sibling->child[!_dir]->color = black;            /*      /  \      /  \     */
	    rotate(_parent, _dir);                            /*     Y3   R5   B1   Y3   */
	    goto END;
	}

	if ( is_red(_sibling->child[_dir]) ){                 /*    X2            X3     */
	    _sibling->child[_dir]->color = _parent->color;    /*   /  \          / \     */
	    _parent->color = black;                           /* (B1)  B4  ->   B2  B4   */
	    rotate(_sibling, !_dir);                          /*      / \      /     \   */
	    rotate(_parent,   _dir);                          /*     R3  B5   B1     B5  */
	    goto END;
	}
	_sibling->color = red;                                /*    X2         (X2)      */
	node = _parent;                                       /*   /  \        /  \      */
    }//while                                                     (B1)  B4  -> B1  R4
END:                                                          /*      /  \        / \    */
    if (_node->child[0] != NULL){                             /*     B3  B5      B3  B5  */
	_node->child[0]->color = _node->color;
	_node->child[0]->idx += _node->idx;
    }
    set_child(_node->parent, _node->child[0], _node->dir);
    free(_node);
    update(tree);
    tree->n--;
    return node_deleted;
}

int _ntree(rbnode node){
    int ans;
    if ( node==NULL )
	return 0;
    if ( is_red(node) && (is_red(node->child[0]) || is_red(node->child[1])) ){
	printf("consecutive red nodes at: ");
	_print(node);
	return -1;
    }
    if ( (ans=_ntree(node->child[0])) != _ntree(node->child[1]) ){
	printf("unequal black nodes at: ");
	_print(node);
	return -1;
    }
    if ( ans==-1 )
	return ans;
    ans += (is_black(node))? 1: 0;
    return ans;
}

